import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import os

class Ic_record():
    def __init__(self, codes,  predict_save=False) -> None:
        self.predict_save = predict_save
        self.codes = codes   
        self.test_ic_list = []
        self.test_rmse_list = []
        self.train_ic_list = []
        self.train_rmse_list = []
        self.eval_rmse_list= []
        self.eval_ic_list= []
        
    def ic_record(self, pred, true, root, times, name='test'):  
        pred_df, true_df, ic_list, rmse_list = self.array_to_df(pred, true, times)
        if name == 'test':
            self.test_rmse_list.append(rmse_list)
            self.test_ic_list.append(ic_list)
        if name == 'eval':
            self.eval_rmse_list.append(rmse_list)
            self.eval_ic_list.append(ic_list)
        if name == 'train':
            self.train_rmse_list.append(rmse_list)
            self.train_ic_list.append(ic_list)          
        print( f'当前{name}IC_mean：{np.nanmean(ic_list)}, IC_std：{np.nanstd(ic_list)},IR:{np.nanmean(ic_list)/np.nanstd(ic_list)}, rmse：{np.nanmean(rmse_list)}')
        if self.predict_save:
            pred_df.to_pickle(os.path.join(root, f'{name}_pred.pkl.gzip'))
            true_df.to_pickle(os.path.join(root, f'{name}_true.pkl.gzip'))
        plot_line_ml(array1=pred_df.iloc[-1,:200].values, array2=true_df.iloc[-1,:200].values,   name=name, root=root) #times=times,
         
    def save_ic(self, root, month_index):
        test_ic_df = ic_re_formate(self.test_ic_list, month_index, self.codes)
        test_ic_df.to_csv(f'{root}/test_ic.csv')
        eval_ic_df = ic_re_formate(self.eval_ic_list, month_index, self.codes)
        eval_ic_df.to_csv(f'{root}/eval_ic.csv')
        train_ic_df = ic_re_formate(self.train_ic_list, month_index, self.codes)
        train_ic_df.to_csv(f'{root}/train_ic.csv')
        # test_rmse_df = ic_re_formate(self.test_rmse_list, month_index, self.codes)
        # train_rmse_df = ic_re_formate(self.train_rmse_list, month_index, self.codes)

    def array_to_df(self, pred, true, eval_times):
        pred = shape_squeeze(pred)
        true = shape_squeeze(true)
        pred_df = pd.DataFrame(pred, index=eval_times, columns=self.codes)#index_time[:len(pred)]
        true_df = pd.DataFrame(true, index=eval_times, columns=self.codes)
        ic_list = []
        cor = calc_rmse(pred.reshape((-1,), order='F'), true.reshape((-1,), order='F'))
        print('整体cor:', cor)
        for i in range(pred.shape[1]):
            cor = calc_corr(pred[:, i], true[:, i])
            ic_list.append(cor)
        rmse_list = []
        for i in range(pred.shape[1]):
            rmse = calc_rmse(pred[:, i], true[:, i])
            rmse_list.append(rmse)
        return pred_df, true_df, ic_list, rmse_list


def ic_re_formate(ic_list, month_index, codes):
    ic = np.array(ic_list)
    ic_df = pd.DataFrame(ic, index=month_index, columns=codes)
    ic_values = ic_df.values
    if len(ic_values.shape)>1:
        mean_id = np.nanmean(ic_values, axis = 0)
        std_id = np.nanstd(ic_values, axis = 0)
        lst = [mean_id, std_id]
        id_df = pd.DataFrame(lst, index=['mean', 'std'], columns=ic_df.columns)
        ic_df = pd.concat([ic_df, id_df], axis=0)
        ic_values_id = ic_df.values 
        mean_id = np.nanmean(ic_values_id, axis = 1)
        std_id = np.nanstd(ic_values_id, axis = 1)
        ic_df['mean'] = mean_id
        ic_df['std'] = std_id
    return ic_df

def calc_rmse(array_a, array_b):
    return np.sqrt(np.nansum(np.square(array_a-array_b)))/len(array_a)

def calc_corr(array_a, array_b):
    if np.nanstd(array_b.tolist())==0:
        print('label无波动，相关系数计算失败')
    if np.nanstd(array_a.tolist())==0:
        print('预测结果无波动，相关系数计算失败')
    value = np.stack([array_a,array_b],axis=1)
    df = pd.DataFrame(value)
    # if self.method == 'spearman':
    #     ic_mat = df.rank().corr(min_periods=1)
    # else:
    ic_mat = df.corr(min_periods=1)
    ic = ic_mat.values[0,1]
    return ic
def shape_squeeze(data):
    shape = data.shape  
    if len(shape)>2:
        data = np.squeeze(data, axis=-1)
    return data

def plot_line_ml(array1, array2, name, root):
    plt.figure(figsize=(30, 10))
    plt.plot(array1,  alpha=0.5, label='pred')
    plt.plot(array2, alpha=0.5, label='true')
    plt.plot(np.ones(len(array1))*0.01, label = '0.01' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*0.02, label = '0.02' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*0.05, label = '0.05' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*-0.01, label = '-0.01' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*-0.02, label = '-0.02' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*-0.05, label = '-0.05' , alpha=0.5, linestyle=':')
    plt.legend()
    save_root = f'{root}/{name}_result.png'
    plt.savefig(save_root)
    plt.close()